# -*- coding: utf-8 -*-
"""
    PBAR_ExamineMarkerValuesVer2.py
    
    Open all the markerValue files for examination.  Makes nice plots.  
    Based on PBAR_ExamineMarkerValues.py.  Got rid of awful dir sub dir structure.

Created on Fri May 02 09:08:03 2014

@author: jkwong
"""

import PBAR_Zspec, PBAR_Cargo
reload(PBAR_Zspec)
reload(PBAR_Cargo)
import numpy as np
import os, copy, glob, cPickle
import matplotlib.pyplot as plt
#from scipy import interpolate
from matplotlib import cm
from scipy import ndimage


# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

# define good/bad detectors
(goodZspecMask, badZspecMask, goodZspecNameList, badZspecNameList, \
        goodZspecIndices, badZspecIndices) = PBAR_Zspec.ZspecDetectorLists2ndSet()

dataPath = r'E:\PBAR\data4\BasicScansStandardWidth'
dataPathPickle = r'E:\PBAR\data4\BasicScansStandardWidthPickle'

plotDir = r'E:\PBAR\data4\BasicScansPlots'
basepathSet2 = r'C:\Users\jkwong\Documents\Work\PBAR\data4\Mar-files'
plotSaveDir = r'E:\PBAR\data4\BasicScansPlots'

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')

# define parameters for standard width images
numberTimeSlices = 2000
preCargoTimeSlices = 70

acqTime = 1/60.
figureSize = (16, 10)

params = {'legend.fontsize': 8}
plt.rcParams.update(params)

#################################
## DEFINE FILE NAMES

markerValuesList = []

filenameMarkerValuesList = []
fullFilenameMarkerValuesList = []

fullFilenameZspecList = []
filenameZspecList = []
fullFilenameCargoList = []
filenameCargoList = []
cargoNumberList = []

datStandardWidthList = []
datCargoStandardWidthList = []

for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):
    print(datasetIndex)
#    if datasetIndex == 1:
#        break
#    if datasetIndex < 24:
#        continue
    
    ######################
    ## READ IN DATA
    
    cargoNumberList.append(datasetDescription['config'][datasetIndex])    
    
    # Read Zspec
#    filenameZspec = '%s-FDFC-All_SW.npy' %datasetDescription['scanID'][datasetIndex]
#    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
#    filenameZspecList.append(filenameZspec)
#    fullFilenameZspecList.append(fullFilenameZspec)
#    print('Loading %s' %fullFilenameZspec)
#    datStandardWidthList.append(np.load(fullFilenameZspec))
    
    filenameZspec = '%s-FDFC-All_SW.dat' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPathPickle, filenameZspec)
    filenameZspecList.append(filenameZspec)
    fullFilenameZspecList.append(fullFilenameZspec)
    print('Loading %s' %fullFilenameZspec)
    energy, datStandardWidth = PBAR_Zspec.ReadZspecBasicScanPickle(fullFilenameZspec)    
    datStandardWidthList.append(np.load(fullFilenameZspec))
    
    
    # Read Cargo
#    filenameCargo = 'PBAR-%s.cargoimageSW.npy' %datasetDescription['dataFile'][datasetIndex]
#    fullFilenameCargo = os.path.join(dataPath, filenameCargo)
#    filenameCargoList.append(filenameCargo)
#    fullFilenameCargoList.append(fullFilenameCargo)
#    print('Loading %s' %fullFilenameCargo)
#    datCargoStandardWidthList.append(np.load(fullFilenameCargo))

    filenameCargo = 'PBAR-%s.cargoimageSW.dat' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPathPickle, filenameCargo)
    filenameCargoList.append(filenameCargo)
    fullFilenameCargoList.append(fullFilenameCargo)    
    print('Loading %s' %fullFilenameCargo)
    with open(fullFilenameCargo, 'rb') as fid:
        datCargoStandardWidth = cPickle.load(fid)
    datCargoStandardWidthList.append(datCargoStandardWidth)    
    
    # Read in marker files
    filenameMarker = filenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')

    filenameMarkerValues = filenameCargo.replace('cargoimageSW', 'cargomarkervaluesSW')
    fullFilenameMarkerValues = fullFilenameCargo.replace('cargoimageSW', 'cargomarkervaluesSW')
    filenameMarkerValuesList.append(filenameMarkerValues)
    fullFilenameMarkerValuesList.append(fullFilenameMarkerValues)
    
# # don't have to read in marker file - i think marker values file already has this stuff
#    # some don't have marker files
#    if os.path.exists(fullFilenameMarker):
##        markers = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
#        markerStandardWidth = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
#    else:
#        markerStandardWidth = []

    # read in the marker value files
    if os.path.exists(fullFilenameMarkerValues):
        with open(fullFilenameMarkerValues, 'rb') as fid:
            markerValuesList.append(cPickle.load(fid))
    else:
        print('Does not exist: %s' %fullFilenameMarkerValues)
        markerValuesList.append([])

discrimBackgroundList = []
featureNameList = ['binMean', 'binKur', 'binSkew', 'count', 'dist0']

cargoCountRange = [0.0, 4e7]
for datasetIndex in xrange(len(datStandardWidthList)):
    print(datasetIndex)
    discrimBackground = {}
    datStandardWidth = datStandardWidthList[datasetIndex]
    datCargoStandardWidth = datCargoStandardWidthList[datasetIndex]
    x = np.reshape(datCargoStandardWidth, datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])
    cut = (x > cargoCountRange[0]) & (x < cargoCountRange[1])
    # interp version
    discrim = PBAR_Zspec.CalculateFeaturesZspecBasic(datStandardWidth, energy)
    for featureName in featureNameList:
        discrimBackground[featureName] = np.reshape(discrim[featureName], datCargoStandardWidth.shape[0] * datCargoStandardWidth.shape[1])[cut]
    discrimBackgroundList.append(discrimBackground)

# load the polynomial
# set 2 - load from file
setNum = 2
fullFilename = os.path.join(basepathSet2, 'zspec%dSet.dat' %setNum)
with open(fullFilename ,'rb') as fid:
    print('Reading %s' %fullFilename)
    temp = cPickle.load(fid)

polyMeanOrderExtendIndex = 22
polySigmaOrderExtendIndex = 13

pfitmean = temp['pfit']['mean']['PbALL']['extend'][polyMeanOrderExtendIndex]
pfitsigma = temp['pfit']['sigma']['PbALL']['extend'][polySigmaOrderExtendIndex]
del temp


# Make nice arrays for the markerValues
#featuresBasicScan[binning type][feature name]

featuresBasicScan = {}
featuresBasicScanMean = {}

boxTypeList = ['center', 'centerBlock', 'block', 'blockMask']

#boxTypeList = ['center']


featureNameList = markerValuesList[0][0]['center']['discrim'].keys()
featureNameList = ['binMean', 'binKur', 'binSkew', 'count', 'dist0', 'cargoCount']


for boxTypeIndex, boxType in enumerate(boxTypeList):
    print(boxType)
    featuresBasicScan[boxType]= {}
    featuresBasicScanMean[boxType] = {}
    
    for featureNameIndex, featureName in enumerate(featureNameList):   # cycle through features
        featuretemp = []
        featuremeantemp = []
        
        materialtemp = []
        materialmeantemp = []

        cargonumbertemp = []
        carognumbermeantemp = []
        
        # build up the feature array by concatenating all the values
        for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets
           
            cargoNumber = cargoNumberList[markerIndex]
            for i, mark in enumerate(marker): # go through markers in a dataset
#                print('cargo config %d, marker %s'  %(cargoNumberList[markerIndex], mark['target']))
#                print(markerIndex, i)
                # center point
                if boxType == 'center':
                    xx = mark['center']['discrim'][featureName]
                    xxMean = mark['center']['discrim'][featureName]
                elif boxType == 'centerBlock':        
                    # center block
                    xx = mark['centerBlock']['discrim'][featureName]
                    xxMean = mark['centerBlock']['discrimMean'][featureName]
                elif boxType == 'block':  
                    # center
                    xx = mark['block']['discrim'][featureName]
                    xxMean = mark['block']['discrimMean'][featureName]
                elif boxType == 'blockMask':
                    # center 
                    xx = mark['blockMask']['discrim'][featureName]
                    xxMean = mark['blockMask']['discrimMean'][featureName]
#                print(xx.shape)
                
#                # skip if the value is nan or inf
#                if np.isnan(yy) or np.isnan(xx) or np.isinf(yy) or np.isinf(xx):
#                    print('Value is nan or inf, skipping')
#                    continue
                
                featuremeantemp.append(xxMean)
                
                try:
                    featuretemp = featuretemp + list(xx)
                except:
                    featuretemp.append(xx)
                
                if (len(mark['target']) > 0):
                    materialmeantemp.append(mark['target'][0])
                    try:
                        materialtemp = materialtemp + len(xx)*[mark['target'][0]]
                    except:
                        materialtemp = materialtemp + [mark['target'][0]]
                else:
                    materialmeantemp.append('_')
                    try:
                        materialtemp = materialtemp + len(xx)*['_']
                    except:
                        materialtemp = materialtemp + ['_']

                carognumbermeantemp.append(cargoNumber)
                try:
                    cargonumbertemp = materialtemp + len(xx)*[cargoNumber]
                except:
                    cargonumbertemp = materialtemp + [cargoNumber]

        featuresBasicScanMean[boxType][featureName] = np.array(featuremeantemp)                    
        featuresBasicScan[boxType][featureName] = np.array(featuretemp)
        
        featuresBasicScanMean[boxType]['material'] = np.array(materialmeantemp)
        featuresBasicScan[boxType]['material'] = np.array(materialtemp)

        featuresBasicScanMean[boxType]['cargoNumber'] = np.array(carognumbermeantemp)
        featuresBasicScan[boxType]['cargoNumber'] = np.array(cargonumbertemp)


# high z masks
markerhighZMask = np.zeros(len(featuresBasicScanMean[boxType]['material'])) -1
for i in xrange(len(markerhighZMask)):
    if (featuresBasicScanMean[boxType]['material'][i].lower() == 'f'):
        markerhighZMask[i] = 0
    elif (featuresBasicScanMean[boxType]['material'][i].lower() == 'p'):
        markerhighZMask[i] = 1
    elif (featuresBasicScanMean[boxType]['material'][i].lower() == 't') or \
        (featuresBasicScanMean[boxType]['material'][i].lower() == 'd') or \
        (featuresBasicScanMean[boxType]['material'][i].lower() == 'w'):
        markerhighZMask[i] = 2

# build up the feature array by concatenating all the values
materialUnique = []
for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets
    cargoNumber = cargoNumberList[markerIndex]
    for i, mark in enumerate(marker): # go through markers in a dataset
        materialUnique.append(mark['target'])


#  ANALYSIS

## old values from 10/2013
#removeOffset = 0
#wBest = np.array([-2.07779891,  1.88095226,  4.65429156])
#pfitmean = np.array([ -3.45931698e-06,   1.80445624e-03,  -3.46939502e-01,  9.03548991e-01])
#pfitsigma = np.array([ -2.90915823e-07,   2.71344856e-04,  -6.81517865e-02, 6.54989380e+00])


# Plots histogram of the center and centerblock values and also the mean value
# Comparing the 'center' and 'centerblock' data points

plt.figure()
plt.grid()
feature1 = 'count'
feature1 = 'dist0'
feature1 = 'cargoCount'
index = 10
markerIndex = 2

pixelGroup = 'block'
hist(markerValuesList[index][markerIndex][pixelGroup]['discrim'][feature1])
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrimMean'][feature1], 0, 'ob', markersize = 15)
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrimMeanSpectrum'][feature1], 0, 'om', markersize = 15)

pixelGroup = 'centerBlock'
hist(markerValuesList[index][markerIndex][pixelGroup]['discrim'][feature1])
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrimMean'][feature1], 0, 'sb', markersize = 15)
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrimMeanSpectrum'][feature1], 0, 'sm', markersize = 15)

pixelGroup = 'centerBlock2'
hist(markerValuesList[index][markerIndex][pixelGroup]['discrim'][feature1])
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrimMean'][feature1], 0, 'xb', markersize = 15)
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrimMeanSpectrum'][feature1], 0, 'xm', markersize = 15)

pixelGroup = 'center'
plt.plot(markerValuesList[index][markerIndex][pixelGroup]['discrim'][feature1], 0, 'or', markersize = 15)

plt.title(filenameZspecList[index] + ', ' + markerValuesList[index][markerIndex]['target'])
plt.legend()
plt.xlabel(feature1)
plt.ylabel('Count')


#  Comparing the features from different pixels groups and mean types
#  
#  If everything is fine, there should be anticorrelation here
#  This will plot a lot of points
#
#center
#centerBlock
#block
#blockMask


plt.figure()
plt.grid()
feature1 = 'count'
#feature1 = 'dist0'
#feature1 = 'binKur'
#feature1 = 'binMean'
#feature1 = 'cargoCount'

#feature2 = 'count'
#feature2 = 'dist0'
#feature2 = 'binKur'
#feature2 = 'binMean'
feature2 = 'cargoCount'


boxType1 = 'center'
boxType1 = 'blockMask'

#boxType2 = 'centerBlock2'
boxType2 = 'blockMask'

meanType1 = 'discrimMean'
#meanType1 = 'discrimMean'
#meanType1 = 'discrimMeanSpectrum'

meanType2 = 'discrimMean'
#meanType2 = 'discrimMeanSpectrum'

removeOffset = 0

markerAlpha = 0.2

for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets          
    for i, mark in enumerate(marker): # go through markers in a dataset
    
#        xx, yy = mark['center']['discrim'][feature1], mark['center']['discrim'][feature2]
        xx = copy.copy(mark[boxType1][meanType1][feature1])
#        yy = copy.copy(mark['centerBlock']['discrimMean'][feature1])
        yy = copy.copy(mark[boxType2][meanType2][feature2])

#        yy = copy.copy(mark['blockMask']['discrimMean'][feature1])
        
        if removeOffset:
            yy = yy - np.polyval(pfitmean, xx)
        try:
            if mark['target'][0] == 'S':
                plt.plot(xx, yy, 'dy', markersize = 12, alpha  = markerAlpha, label = mark['target'] +' Center')
            elif mark['target'][0] == 'W':
                plt.plot(xx, yy, 'vb', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            elif mark['target'][0] == 'D':
                plt.plot(xx, yy, 'sb', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            elif mark['target'][0] == 'P':
                plt.plot(xx, yy, '*c', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            elif mark['target'][0] == 'F':# low density stuff
                plt.plot(xx, yy, 'or', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            else:
                plt.plot(xx, yy, 'om', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
        except:
            plt.plot(xx, yy, 'om', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            
plt.title('Discriminant vs Counts, Cargos 1-48')
plt.legend()
plt.xlabel(boxType1 + ', ' + meanType1 + ', ' +feature1)
plt.ylabel(boxType2 + ', '  + meanType2 + ', ' + feature1)
#plt.xlim((0, 500))
#plt.ylim((-100, 100))


for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets
    for i, mark in enumerate(marker): # go through markers in a dataset
        if mark['target'] == 'Center':
            print(markerIndex, i)



# scatter plot of the discriminant vs zspec count

plt.figure()
plt.grid()
plotAlpha = 0.3
markSize = 12
labelPoints = 1
markerAlpha = 0.2
textAlpha = 0.5

# pixel group
boxType = 'center'
#boxType = 'centerBlock'
#boxType = 'centerBlock2'
#boxType = 'block'
#boxType = 'blockMask'

#meanType = 'discrimMean'
meanType = 'discrimMeanSpectrum'


if boxType == 'center':
    meanType = 'discrim'

removeOffset = 0

#
#feature1 = 'count'

feature1 = 'cargoCount'

feature2 = 'dist0'
#feature2 = 'binKur'
#feature2 = 'binSkew'
#feature2 = 'binMean'


cargoCountRange = np.array([0.0, 0.4e8])

numberDataPoints = 0

for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets          
    for i, mark in enumerate(marker): # go through markers in a dataset
    
    
        # center point
        if boxType == 'center':
            xx = mark['center']['discrim'][feature1]
            yy = mark['center']['discrim'][feature2]
            inRange = ( mark['center']['discrim']['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['center']['discrim']['cargoCount'] <= cargoCountRange[1] )

        elif boxType == 'centerBlock':
            # center block
            xx = mark['centerBlock'][meanType][feature1]
            yy = mark['centerBlock'][meanType][feature2]        
            inRange = ( mark['centerBlock'][meanType]['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['centerBlock'][meanType]['cargoCount'] <= cargoCountRange[1] )
        elif boxType == 'centerBlock2':
            # center block
            xx = mark[boxType][meanType][feature1]
            yy = mark[boxType][meanType][feature2]        
            inRange = ( mark[boxType][meanType]['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark[boxType][meanType]['cargoCount'] <= cargoCountRange[1] )
        elif boxType == 'block':  
            # center 
            xx = mark['block'][meanType][feature1]
            yy = mark['block'][meanType][feature2]        
            inRange = ( mark['block'][meanType]['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['block'][meanType]['cargoCount'] <= cargoCountRange[1] )
        elif boxType == 'blockMask':  
            # center 
            xx = mark['blockMask'][meanType][feature1]
            yy = mark['blockMask'][meanType][feature2]        
            inRange = ( mark['blockMask'][meanType]['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['blockMask'][meanType]['cargoCount'] <= cargoCountRange[1] )            

#        if ~inRange:
#            print('Skipping')
#            continue
        
        if removeOffset:
            yy = yy - np.polyval(pfitmean, xx)
            

        print('cargo config %d, marker %s, inRange %d, xx = %3.3f, yy == %3.3f'  %(cargoNumberList[markerIndex], mark['target'], inRange, xx, yy))            
            
        
        # skip if the value is nan or inf
        if np.isnan(yy) or np.isnan(xx) or np.isinf(yy) or np.isinf(xx):
            print('Value is nan or inf, skipping')
            continue
        
        try:
            if mark['target'][0] == 'S':
                markerType = 'dy'
            elif mark['target'][0] == 'W':
                markerType = 'vb'
            elif mark['target'][0] == 'D':
                markerType = 'sb'
            elif mark['target'][0] == 'P':
                markerType = '*c'
            elif mark['target'][0] == 'F':# low density stuff
                markerType = 'or'
            elif mark['target'].lower()[0] == 't':# threat
                markerType = '^b'
            else:
                markerType = 'om'
        except:
            markerType = 'om'
        
        numberDataPoints += 1
        plt.plot(xx, yy, markerType, markersize = markSize, alpha = plotAlpha, label = '%s, %s, %s, ' %(mark['target'], boxType, meanType) )
        
#        if cargoNumberList[markerIndex] == 29:
#            if mark['target'][0] == 'F':
#                plt.plot(xx, yy, markerType, markersize = 18, alpha = markerAlpha, label = '%s, %s, %s, ' %(mark['target'], boxType, meanType) )
#                if ~np.isnan(yy) and ~np.isnan(xx):
#                    plt.text(xx, yy, mark['target'], color = 'g', fontsize = 16)
#                print(xx, yy)
#                
        # print number next to data point
        if labelPoints:
            plt.text(xx, yy, cargoNumberList[markerIndex], fontsize = 10, color = markerType[-1], alpha = textAlpha)
        
plt.title('Discriminant vs Counts, %s, Remove offset: %d, box type: %s,  mean type: %s' %(boxType, removeOffset,boxType, meanType))
plt.legend(prop={'size':6})
plt.xlabel(feature1, fontsize = 16)
plt.ylabel(feature2, fontsize = 16)

if feature1 == 'count':
    plt.xlim((0, 500))

print(numberDataPoints)
if feature2 == 'dist0':
    plt.ylim((-40, 40))
elif feature2 == 'binKur':
    plt.ylim((0, 15))
elif feature2 == 'binSkew':
    plt.ylim((0, 4))




# scatter plot of the discriminant vs zspec count with new data structure

plotAlpha = 0.3
markSize = 12
labelPoints = 1
markerAlpha = 0.2
textAlpha = 0.5

boxType = 'center'
#boxType = 'centerBlock'
#boxType = 'block'
#boxType = 'blockMask'

plt.figure()
plt.grid()
feature1 = 'count'
#feature1 = 'cargoCount'
#feature2 = 'binKur'
#feature2 = 'binSkew'
#feature2 = 'binMean'
feature2 = 'dist0'

# Plot the background points

index = 9

if feature2 == 'binMean':
    cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
        (discrimBackgroundList[index][feature2] < 70) & (discrimBackgroundList[index][feature2] > 10)
elif feature2 == 'binSkew':
    cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
        (discrimBackgroundList[index][feature2] > 0) & (discrimBackgroundList[index][feature2] < 3)    
    
pfit = np.polyfit(discrimBackgroundList[index][feature1][cut], discrimBackgroundList[index][feature2][cut], 10)
plt.plot(discrimBackgroundList[index][feature1], discrimBackgroundList[index][feature2], '.m', markersize = 8, alpha = 0.2, label = 'Background Points')

xx = featuresBasicScanMean[boxType][feature1]
yy = featuresBasicScanMean[boxType][feature2]
material = featuresBasicScanMean[boxType]['material']

cut = markerhighZMask == 2
plt.plot(xx[cut], yy[cut], 'ob', markersize = 12, alpha  = 1.0, label = 'box %s, High Z' %(boxType))
#cut = markerhighZMask == 1
#plt.plot(xx[cut], yy[cut], 'om', markersize = 12, alpha  = 1.0, label = 'box %s, Pb' %(boxType))
cut = markerhighZMask == 0
plt.plot(xx[cut], yy[cut], 'or', markersize = 12, alpha  = 1.0, label = 'box %s, Low Z' %(boxType))

xx = np.arange(400)
yy = np.polyval(pfit, xx)
plt.plot(xx, yy)

plt.title('zspec counts vs cargo counts, Cargos 1-39')
plt.legend()
plt.xlabel(feature1)
plt.ylabel(feature2)
#plt.xlim((0, 500))
#plt.ylim((-100, 100))
plt.legend(prop={'size':16})
plt.xlim((0,400))


# scatter plot of the discriminant vs zspec count with new data structure
# offset removed

plotAlpha = 0.3
markSize = 12
labelPoints = 1
markerAlpha = 0.2
textAlpha = 0.5

removeOffset = 1

boxType = 'center'
#boxType = 'centerBlock'
#boxType = 'block'
#boxType = 'blockMask'

plt.figure()
plt.grid()
feature1 = 'count'
#feature1 = 'cargoCount'
#feature2 = 'binKur'
#feature2 = 'binSkew'
#feature2 = 'binMean'
feature2 = 'dist0'

# Plot the background points

index = 9

if feature2 == 'binMean':
    cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
        (discrimBackgroundList[index][feature2] < 70) & (discrimBackgroundList[index][feature2] > 10)
elif feature2 == 'binSkew':
    cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
        (discrimBackgroundList[index][feature2] > 0) & (discrimBackgroundList[index][feature2] < 3)    

# plot background points
#pfit = np.polyfit(discrimBackgroundList[index][feature1][cut], discrimBackgroundList[index][feature2][cut], 10)
xx = discrimBackgroundList[index][feature1]
yy = discrimBackgroundList[index][feature2]
if removeOffset:
    yy = yy - np.polyval(pfitmean, xx)    
plt.plot(xx, yy, '.m', markersize = 8, alpha = 0.2, label = 'Background Points')

xx = featuresBasicScanMean[boxType][feature1]
yy = featuresBasicScanMean[boxType][feature2]
material = featuresBasicScanMean[boxType]['material']

# plot marker points
if removeOffset:
    yy = yy - np.polyval(pfitmean, xx)

cut = markerhighZMask == 2
plt.plot(xx[cut], yy[cut], 'ob', markersize = 12, alpha  = 1.0, label = 'box %s, High Z' %(boxType))
#cut = markerhighZMask == 1
#plt.plot(xx[cut], yy[cut], 'om', markersize = 12, alpha  = 1.0, label = 'box %s, Pb' %(boxType))
cut = markerhighZMask == 0
plt.plot(xx[cut], yy[cut], 'or', markersize = 12, alpha  = 1.0, label = 'box %s, Low Z' %(boxType))

#xx = np.arange(400)
#yy = np.polyval(pfit, xx)
#plt.plot(xx, yy)

plt.title('zspec counts vs cargo counts, Cargos 1-39')
plt.legend()
plt.xlabel(feature1)
plt.ylabel(feature2)
#plt.xlim((0, 500))
#plt.ylim((-100, 100))
plt.legend(prop={'size':16})
plt.xlim((0,400))
if feature2 == 'dist0':
    plt.ylim((-40, 40))




# plot a couple backgrounds

plt.figure()
plt.grid()
feature1 = 'count'
#feature1 = 'cargoCount'
#feature2 = 'binKur'
feature2 = 'binSkew'
#feature2 = 'binMean'
#feature2 = 'dist0'
for index in xrange(10):
    if feature2 == 'binMean':
        cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
            (discrimBackgroundList[index][feature2] < 70) & (discrimBackgroundList[index][feature2] > 10)
    elif feature2 == 'binSkew':
        cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
            (discrimBackgroundList[index][feature2] > 0) & (discrimBackgroundList[index][feature2] < 3)
    pfit = np.polyfit(discrimBackgroundList[index][feature1][cut], discrimBackgroundList[index][feature2][cut], 10)
    plt.plot(discrimBackgroundList[index][feature1], discrimBackgroundList[index][feature2], '.', alpha = 0.2, label = 'Background Points')
    xx = np.arange(400)
    yy = np.polyval(pfit, xx)
    plt.plot(xx, yy)

plt.axis((0, 400, 0, 70))



# plot a couple backgrounds

feature1 = 'count'
#feature1 = 'cargoCount'
#feature2 = 'binKur'
feature2 = 'binSkew'
#feature2 = 'binMean'
#feature2 = 'dist0'
for index in xrange(10):
    plt.figure()
    plt.grid()
    cut = (discrimBackgroundList[index][feature1] > 0) & (discrimBackgroundList[index][feature1] < 300) & \
        (discrimBackgroundList[index][feature2] < 70) & (discrimBackgroundList[index][feature2] > 10)
    pfit = np.polyfit(discrimBackgroundList[index][feature1][cut], discrimBackgroundList[index][feature2][cut], 10)
    plt.plot(discrimBackgroundList[index][feature1], discrimBackgroundList[index][feature2], '.', alpha = 0.2, label = 'Background Points')
    xx = np.arange(400)
    yy = np.polyval(pfit, xx)
    plt.plot(xx, yy)
    plt.title(datasetDescription['scanID'][index])
    plt.axis((0, 400, 0, 70))
    plt.xlabel(feature1)
    plt.ylabel(feature2)

